import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs

class KMeans:
    def __init__(self, k=3, max_iters=100, random_state=42):
        self.k = k
        self.max_iters = max_iters
        self.random_state = random_state
        
    def fit(self, X):
        np.random.seed(self.random_state)
        n_samples, n_features = X.shape
        
        # 随机初始化中心点
        self.centers = X[np.random.choice(n_samples, self.k, replace=False)]
        
        # 存储历史中心点用于可视化
        self.history_centers = [self.centers.copy()]
        self.history_labels = []
        
        for iteration in range(self.max_iters):
            # 分配标签
            distances = np.linalg.norm(X[:, np.newaxis] - self.centers, axis=2)
            labels = np.argmin(distances, axis=1)
            
            # 存储历史标签
            self.history_labels.append(labels.copy())
            
            # 更新中心点
            new_centers = np.array([X[labels == i].mean(axis=0) for i in range(self.k)])
            
            # 检查收敛
            if np.allclose(self.centers, new_centers):
                print(f"在第 {iteration + 1} 次迭代后收敛")
                break
                
            self.centers = new_centers
            self.history_centers.append(self.centers.copy())
        
        self.labels_ = labels
        return self
    
    def predict(self, X):
        distances = np.linalg.norm(X[:, np.newaxis] - self.centers, axis=2)
        return np.argmin(distances, axis=1)

def generate_sample_data():
    """生成示例数据"""
    X, y_true = make_blobs(n_samples=300, centers=3, n_features=2, 
                          random_state=42, cluster_std=1.5)
    return X, y_true

def visualize_kmeans_process(X, kmeans, y_true):
    """可视化K-means过程"""
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    axes = axes.ravel()
    
    # 选择几个关键迭代步骤进行可视化
    iterations_to_show = [0, 1, 2, 3, 4, len(kmeans.history_labels)-1]
    
    for idx, iter_num in enumerate(iterations_to_show):
        if iter_num < len(kmeans.history_labels):
            labels = kmeans.history_labels[iter_num]
            centers = kmeans.history_centers[iter_num]
            
            # 绘制数据点
            scatter = axes[idx].scatter(X[:, 0], X[:, 1], c=labels, cmap='viridis', 
                                      alpha=0.6, s=50)
            # 绘制中心点
            axes[idx].scatter(centers[:, 0], centers[:, 1], c='red', marker='x', 
                            s=200, linewidths=3, label='中心点')
            
            axes[idx].set_title(f'迭代 {iter_num + 1}')
            axes[idx].set_xlabel('特征 1')
            axes[idx].set_ylabel('特征 2')
            axes[idx].legend()
            axes[idx].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('kmeans_process.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # 绘制最终结果对比
    plt.figure(figsize=(12, 5))
    
    # 真实分布
    plt.subplot(1, 2, 1)
    plt.scatter(X[:, 0], X[:, 1], c=y_true, cmap='viridis', alpha=0.6, s=50)
    plt.title('真实聚类分布')
    plt.xlabel('特征 1')
    plt.ylabel('特征 2')
    plt.grid(True, alpha=0.3)
    
    # K-means结果
    plt.subplot(1, 2, 2)
    plt.scatter(X[:, 0], X[:, 1], c=kmeans.labels_, cmap='viridis', alpha=0.6, s=50)
    plt.scatter(kmeans.centers[:, 0], kmeans.centers[:, 1], c='red', marker='x', 
               s=200, linewidths=3, label='最终中心点')
    plt.title('K-means聚类结果')
    plt.xlabel('特征 1')
    plt.ylabel('特征 2')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('kmeans_final_results.png', dpi=300, bbox_inches='tight')
    plt.show()

def run_kmeans():
    print("=" * 50)
    print("K-means算法实验")
    print("=" * 50)
    
    # 生成数据
    print("正在生成示例数据...")
    X, y_true = generate_sample_data()
    print(f"数据形状: {X.shape}")
    print(f"真实类别数: {len(np.unique(y_true))}")
    
    # 运行K-means
    print("正在运行K-means算法...")
    kmeans = KMeans(k=3, random_state=42)
    kmeans.fit(X)
    
    # 计算准确率（调整标签以匹配真实标签）
    from sklearn.metrics import adjusted_rand_score
    ari = adjusted_rand_score(y_true, kmeans.labels_)
    print(f"调整兰德指数: {ari:.4f}")
    
    # 输出中心点
    print("\n最终中心点坐标:")
    for i, center in enumerate(kmeans.centers):
        print(f"  中心点 {i}: ({center[0]:.2f}, {center[1]:.2f})")
    
    # 可视化结果
    visualize_kmeans_process(X, kmeans, y_true)
    
    return ari

if __name__ == "__main__":
    run_kmeans()